use strict;
use warnings;
$| = 1;
use DBI;
use Cwd;

####################################################################################
# 				newGOVer.pl May 2015
# Purpose: Download and install GO MySQL tables, and cross-reference to Uniprot
#
# You need a reasonably fast computer. If there are many UniProt entries, this can
# take over 12 hours to run.
# 
# To use:
# 1. Set $GOurl to the current GO distribution, or whichever is desired
# 2. Set $UPDir to the directory containing the Uniprot taxonomic directories
#	(In other words, relevant .dat files should be in subdirectories of $UPDir;
#	if you used newUPVer.pl, it will look like "Uniprot_<date>")
# 3. Set $GODB to the desired name for the GO database 
# 4. Copy the HOSTS.cfg from your TCW directory, as it contains the database access information
# 
# Result:
# 1. GO distribution downloaded
# 2. GO MySQL database is created with name $GODB, GO tables loaded, and PAVE_Uniprot table created
# 3. Uniprot .dat files scanned, and GO, Pfam, Kegg, EC cross-references loaded to PAVE_Uniprot
# Note, this database will further be updated the first time you run TCW.
# 
######################################################################################
	
# edit these variables (except if your running demo, and this is run from the DBfasta directory)
my $GOurl = "http://archive.geneontology.org/full/2012-12-01/go_201212-termdb-tables.tar.gz"; 
my $UPDir = "projects/DBfasta/demo4";
my $GODB = "go_demo";
my $gzipDATs = 0; 	# gzip .dat files when done. Or they can be removed.

# The MySQl userid and password from the FIRST entry of HOSTS.cfg will be used! 
# Alternately, comment these lines and edit the database info below. 
my $HOSTScfg = "HOSTS.cfg";
my %Params;
readHostsFile($HOSTScfg,\%Params);

my $dbhost = $Params{DB_host}; # host for GO database (should be same as PAVE project to be annotated)
my $dbuser = $Params{DB_user}; # user having write access
my $dbpass = $Params{DB_password}; # password 

my $GODir = "GO_temp"; 
my $downloadGO = 1;	# Download the GO distribution
my $loadGO = 1;		# Load GO distribution to DB
my $loadUP = 1;		# Load Uniprot names to PAVE_Uniprot table in GO DB
my $loadDAT = 1;	# Load GO (also Kegg,Pfam) assignments from Uniprot .dat files

if (not -d $UPDir)
{
	die "Uniprot directory $UPDir not found. Edit newGOver.pl to define the UPDir variable.\n";
}

my $dbh = DBI->connect("dbi:mysql:mysql;host=$dbhost;", $dbuser, $dbpass) or die "db connect fail";
my $sth;
my %allUP;

my @subDirs;
foreach my $d (<$UPDir/*>)
{
	if (-d $d)
    {
    	push @subDirs, $d;
    }	
}

if ($downloadGO == 1)
{
	if (-d $GODir)
    {
    	print "The GO temporary directory $GODir already exists.\n";
        print "Either delete it or edit the variable GODir in this script, specify a\n";
        print "different directory to use for downloading the GO distribution.\n";
        exit(0);
    }
    else
   	{
    	mkdir $GODir;
    }
    if (not -d $GODir)
    {
    	die "Failed to create GO temporary dir $GODir\n";
    }	
}

if ($loadGO)
{
	# Get the GO and load
    
	check_curl();
    
	my $curdir = cwd();
	chdir $GODir;

	db_exec("show databases like '$GODB'");
	if ($sth->fetchrow_array())
	{
		if (!yesno("Database $GODB already exists; delete it?"))
		{
			exit;
		}
		else
		{
			db_exec("drop database $GODB");
		}
	}

	if ($downloadGO == 1)
    {
		execute("curl -O  $GOurl");
		$GOurl =~ s/.*\///;
		execute("tar xf $GOurl");
		execute("mv *tables/* .");
		unless (-e "term.txt") { # May15 add this check
			print "File does not exist: $GOurl\n";
			print "Check GO site for the latest and change the GO URL variable in this script\n";
			exit;
		}
   	}
    else
    {
    	my @gzs = <*.tar.gz>;
        if (scalar(@gzs) > 1)
        {
        	print "There is more than one .tar.gz file in the $GODB directory.\n";
            print "Please clear this directory and start again.\n";
            exit(0);
        }
        elsif (scalar(@gzs) == 1)
        {
        	my $gz = shift @gzs;
			execute("tar xf $gz");
			execute("cp *tables/* .");
            execute("gzip $gz");
       	}

    }	
	execute("mysqladmin -h $dbhost -u $dbuser -p$dbpass create $GODB");
	execute("cat *.sql | mysql -h $dbhost -u $dbuser -p$dbpass  $GODB");
	execute("mysqlimport -h $dbhost -u $dbuser -L $GODB -p$dbpass *.txt");
    
    chdir($curdir);
}

if ($loadUP)
{
	# Load the Uniprot IDs

	$dbh = DBI->connect("dbi:mysql:$GODB;host=$dbhost;", $dbuser, $dbpass);

	db_exec("CREATE TABLE PAVE_Uniprot (
	                     UPindex bigint unsigned NOT NULL AUTO_INCREMENT PRIMARY KEY,
	                     UPid varchar(25), 
	                     acc varchar(15), 
	                     org varchar(10), 
	                     src ENUM('tr', 'sp'),
	                     updesc text,
	                     prot_seq text,
	                     go text,
	                     pfam tinytext,
                         kegg tinytext,
                         ec tinytext,
	                     unique (UPid)) ");
	$sth = $dbh->prepare("insert ignore into PAVE_Uniprot (upid,acc,org,src,updesc,prot_seq) values(?,?,?,?,?,?)");

	print "Scanning Uniprot fasta to fill PAVE_Uniprot table\n";
	my $firstLine = "";
	my $seq;
	my %allUP;
	my $insertCnt = 0;
    foreach my $dir (@subDirs)
    {  
	    foreach my $file (<$dir/*.fasta>)
	    {
		    print "Loading $file\n";
		    open F, $file or die $file;
	        while (my $line = <F>)
	        {
	    	    if ($line =~ />/)
	            {
	        	    if ($firstLine ne "")
	                {
	            	    handle_up_seq($firstLine,$seq,$sth,$dbh,\%allUP);
	                    $insertCnt++;
	                    if ($insertCnt % 10000 == 0)
	                    {
	                	    print "$insertCnt              \r";
	                    }
	                } 
	                chomp $line;
	                $firstLine = $line;     
	                $seq = "";      
	            }
	            else
	            {
	        	    chomp $line;
	                $line =~ s/\s+//g;
	                $seq .= $line;
	            }
	        }
	        close F;
	    }
    }
	if ($firstLine ne "")
	{
	    handle_up_seq($firstLine,$seq,$sth,$dbh,\%allUP);
	    $insertCnt++;
	}
	$sth->finish;
    
    $sth = $dbh->prepare("update PAVE_Uniprot set go='',pfam='',kegg=''");
    $sth->execute();
    $sth = $dbh->prepare("update PAVE_Uniprot set updesc='' where updesc is null");
    $sth->execute();

	print "Added $insertCnt Uniprots\n";   
}

if ($loadDAT)
{
	print "Scanning Uniprot DATs to fill $GODB table\n";

	$dbh = DBI->connect("dbi:mysql:$GODB;host=$dbhost;", $dbuser, $dbpass);
    my $sth = $dbh->prepare("update PAVE_Uniprot set go=?,pfam=?,kegg=?,ec=? where UPid=?");
    foreach my $dir (@subDirs)
    {  
	    foreach my $dat (<$dir/*.dat>,<$dir/*.dat.gz>)
        {
        	if ($dat =~ /\.gz$/)
            {
            	execute("gunzip $dat");
                $dat =~ s/\.gz$//;
            }
        	my $upnum = 0;
            my %updata;
            open F, $dat or die $dat;
            while (my $line = <F>)
            {
	            if ($line =~ /^ID\s+(\S+)/)
                {
                    $upnum++;
                    updateUP(\%updata, $sth);
                    %updata = ();
                    $updata{ID} = $1;
                	print "$dat\t$upnum\t".$updata{ID}."\r" if $upnum % 1000 == 0;
               	}
                elsif ($line =~ /^DR\s+GO.*(GO:\d+)/)
                {
                	if ($line =~ /^DR\s+GO.*(GO:\d+).*;\s*(\w\w\w?):[^:;]+$/)
                    {
                		my $go = $1;
                    	my $code = $2;
                    	$updata{GO} .= (defined $updata{GO} ? ";$go:$code" : "$go:$code");
                  	}
                    else
                    {
                    	print "No EC: $line";
                    }
                }
                elsif ($line =~ /^DR\s+Pfam.*(PF\d+)/)
                {
                	my $pf = $1;
                    $updata{PF} .= (defined $updata{PF} ? ";$pf" : "$pf");
                }
                elsif ($line =~ /^DR\s+KEGG;\s+(\S+);/)
                {
                	my $kg = $1;
                    $updata{KG} .= (defined $updata{KG} ? ";$kg" : "$kg");
                }
                elsif ($line =~ /^DE\s+EC\s*=\s*([\d.]+);/)
                {
                	my $ec = $1;
                    $updata{EC} .= (defined $updata{EC} ? ";$ec" : "$ec");
                }

            }
            close F;
           	updateUP(\%updata, $sth);
            print "$dat done                                          \n";
        }
	}
}

if ($downloadGO==1 and -d $GODir)
{
	system("rm -Rf $GODir");
}	
if ($gzipDATs) { # May15 moved gzipping to end
	foreach my $dir (@subDirs)
    {  
	    foreach my $dat (<$dir/*.dat>)
        {
        	execute("gzip $dat");
        }
    }
}

###############################################################

sub updateUP
{
	my $p = shift;
    my $sth = shift;
    
    my $id = $p->{ID};
    my $go = (defined $p->{GO} ? $p->{GO} : "");
    my $pf = (defined $p->{PF} ? $p->{PF} : "");
    my $kg = (defined $p->{KG} ? $p->{KG} : "");
    my $ec = (defined $p->{EC} ? $p->{EC} : "");
    
    $sth->execute($go,$pf,$kg,$ec,$id);
    
}

###############################################################

sub handle_up_seq
{
	my $firstLine = shift;
    my $seq = shift;
    my $sth = shift;
    my $dbh = shift;
    my $allUP = shift;

    if ($firstLine =~ />(\S+)\s(\S.*)/)
    {
        my $name = $1;
        my $desc = $2;
                
        my ($src,$acc,$id) = split /\|/, $name;
        if (!$acc || !$id || !($src eq "sp" || $src eq "tr"))
        {
            die "weird desc $firstLine\n";
        }
        my (undef,$org) = split /\_/, $id;
        if (not defined $org)
        {
            die "$name has no org\n";
        }  
        $allUP{$id} = 1;
        $sth->execute($id,$acc,$org,$src,$desc,$seq);
   	}
    else
    {
    	die "can't parse $firstLine\n";
    }
}

###################################################################################

# Takes cmd as first arg.
# Second arg, if present, causes check mode if 1 or execute mode if 0.
sub execute {
	my $cmd = shift;

	print "executing: $cmd\n";
	
	system($cmd);
}

###################################################################################

sub db_exec
{
	my $sql = shift;
	$sth = $dbh->prepare($sql);
	$sth->execute();
}

###################################################################################

sub yesno 
{
	my $query = shift;

	print "$query (y/n) ";
	my $yn = <STDIN>;
	chomp $yn;
	if ( $yn eq "y" ) {
		return 1;
	}
	return 0;
}


#######################################################

sub readHostsFile
{
	my $file = shift;
   	my $pParm = shift;
    unless (-e $file) {
    	print "Error: the HOSTS.cfg must be copied to this directory for the MySQL userid and password\n";
    	exit;
    }
   	open IF, "$file" or die("Failed to locate $file");
    while (<IF>) {
        my ($line) = split /\#/;
        next if (!$line);
        chomp $line;
        next if not $line =~ /=/;
        my ($param, $val) = split /\s*=\s*/, $line;
        $param =~ s/^\s+//;
        $val =~ s/^\s+//;
        $param =~ s/\s+$//;
        $val =~ s/\s+$//;
        $param =~ s/^PAVE_/DB_/;
        return if defined $pParm->{$param};
        $pParm->{$param} = ( $val ? $val : "");

    }
}

##################################################################

sub check_curl
{
        my $curl = `which curl`;
        if (not $curl =~ /^\//)
        {
    	    print "Please install the curl utility for downloading Uniprot files\n";
    	    print "For example, 'sudo yum install curl'\n";
            exit(1);
        }
}

